import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import torch
from PIL import Image
from torchvision import transforms

import numpy as np
import matplotlib.pyplot as plt
import copy
import math
import cv2
from deepDeformOrt import *
import sys

def transStr2Float(strIn, isGrad=0):
    floatList = []
    content = strIn.split(" ")
    for i in range(len(content)-1):
        tmp = float(content[i])
        floatList.append(tmp)
    floatList = torch.Tensor(floatList)
    if 1 == isGrad:
        floatList.requires_grad_(True)
    return floatList

def readMnt(mntName):
    PI = 3.1415926
    dataList = []
    with open("data/" + mntName, 'r') as f:
        for line in f.readlines():
            #print(line)
            dataList.append(line)
    num = len(dataList)
    floatXList = []
    floatYList = []
    floatOrtXList = []
    floatOrtYList = []
    for i in range(num):
        content = dataList[i].split(" ")
        floatXList.append(float(content[1]))
        floatYList.append(float(content[2]))
        tmp = 90-float(content[3])
        tmp = tmp*PI/180.0
        floatOrtXList.append(math.cos(tmp))
        floatOrtYList.append(math.sin(tmp))
    pX = torch.Tensor(floatXList)
    pY = torch.Tensor(floatYList)
    pOrtX = torch.Tensor(floatOrtXList)
    pOrtY = torch.Tensor(floatOrtYList)

    return [pX, pY, pOrtX, pOrtY]
    #return floatXList, floatYList, floatOrtXList, floatOrtYList

def GetPairPoint(pointList1, pointList2, pairName):
    pX1Tmp = pointList1[0]
    pY1Tmp = pointList1[1]
    pOrtX1Tmp = pointList1[2]
    pOrtY1Tmp = pointList1[3]

    pX2Tmp = pointList2[0]
    pY2Tmp = pointList2[1]
    pOrtX2Tmp = pointList2[2]
    pOrtY2Tmp = pointList2[3]

    floatXList1 = []
    floatYList1 = []
    floatOrtXList1 = []
    floatOrtYList1 = []

    floatXList2 = []
    floatYList2 = []
    floatOrtXList2 = []
    floatOrtYList2 = []

    with open("data/" + pairName, 'r') as f:
        for line in f.readlines():
            content = line.split(" ")
            index1 = int(content[0])
            index2 = int(content[1])
            floatXList1.append(pX1Tmp[index1])
            floatYList1.append(pY1Tmp[index1])
            floatOrtXList1.append(pOrtX1Tmp[index1])
            floatOrtYList1.append(pOrtY1Tmp[index1])

            floatXList2.append(pX2Tmp[index2])
            floatYList2.append(pY2Tmp[index2])
            floatOrtXList2.append(pOrtX2Tmp[index2])
            floatOrtYList2.append(pOrtY2Tmp[index2])
    pX1 = torch.Tensor(floatXList1)
    pX1 = pX1.unsqueeze(1)
    pX1.requires_grad_(True)
    pY1 = torch.Tensor(floatYList1)
    pY1 = pY1.unsqueeze(1)
    pY1.requires_grad_(True)
    pOrtX1 = torch.Tensor(floatOrtXList1)
    pOrtY1 = torch.Tensor(floatOrtYList1)

    pX2 = torch.Tensor(floatXList2)
    pX2 = pX2.unsqueeze(1)
    pX2.requires_grad_(True)
    pY2 = torch.Tensor(floatYList2)
    pY2 = pY2.unsqueeze(1)
    pY2.requires_grad_(True)
    pOrtX2 = torch.Tensor(floatOrtXList2)
    pOrtY2 = torch.Tensor(floatOrtYList2)

    return [pX1, pY1, pOrtX1, pOrtY1], [pX2, pY2, pOrtX2, pOrtY2]

def LoadImageData(name1, name2):
    pointAllList1 = readMnt(name1 + ".txt")
    pointAllList2 = readMnt(name2 + ".txt")
    pairName = name1 + "_" + name2 + ".txt"
    pointPairList1, pointPairList2 = GetPairPoint(pointAllList1, pointAllList2, pairName)
    # pX1 = pointList1[0]
    # pY1 = pointList1[1]
    # pOrtX1 = pointList1[2]
    # pOrtY1 = pointList1[3]
    #
    # pX2 = pointList2[0]
    # pY2 = pointList2[1]
    # pOrtX2 = pointList2[2]
    # pOrtY2 = pointList2[3]

    imageName1 = "data/" + name1 + ".bmp"
    imageName2 = "data/" + name2 + ".bmp"
    image1 = cv2.imread(imageName1, 0)
    image2 = cv2.imread(imageName2, 0)

    #return [pAllX1, pAllY1, pAllOrtX1, pAllOrtY1, pAllX2, pAllY2, pAllOrtX2, pAllOrtY2], [pX1, pY1, pOrtX1, pOrtY1, pX2, pY2, pOrtX2, pOrtY2], image1, image2
    return pointAllList1+pointAllList2, pointPairList1+pointPairList2, image1, image2

def getImageTransPoints(width, height, device):
    xList = []
    yList = []
    for i in range(height):
        for j in range(width):
            xList.append(float(j))
            yList.append(float(i))
    pImageX = torch.Tensor(xList)
    pImageY = torch.Tensor(yList)
    pImageX = pImageX.unsqueeze(1)
    pImageY = pImageY.unsqueeze(1)

    pImageX = pImageX.to(device)
    pImageY = pImageY.to(device)
    return pImageX, pImageY

def GetTransImagePixel(width, height, pTransImageX, pTransImageY, image1, image2):
    transImage = np.zeros(height*width, dtype=np.uint8)
    overlapNum = 0
    mse = 0
    for i in range(len(pTransImageX)):
        y1 = i//width
        x1 = i%width
        x = round(float(pTransImageX[i]))
        y = round(float(pTransImageY[i]))
        if 0<=x and x<width and 0<=y and y<height:
            transImage[i] = image2[y, x]
            mse = mse + (image1[y1, x1]-image2[y, x])*(image1[y1, x1]-image2[y, x])
            overlapNum = overlapNum + 1
    transImage = transImage.reshape(height, width)
    if 0 < overlapNum:
        mse = mse/overlapNum
    return transImage, mse

def PlotDirCircleCV(pX, pY, pOrtX, pOrtY, image):
    r1 = 16
    point_size = 5
    color = (255, 255, 0)  # BGR
    thickness = 2  # 0 、4、8

    for i in range(len(pX)):
        cv2.circle(image, (int(pX[i]), int(pY[i])), point_size, color, thickness)
        X = [pX[i], pX[i] + r1 * pOrtX[i]+0.5]
        Y = [pY[i], (pY[i] + r1 * pOrtY[i])+0.5]
        cv2.line(image,(int(X[0]),int(Y[0])),(int(X[1]),int(Y[1])),color,thickness,cv2.LINE_4)

def Plot2DLineCV(pX1, pY1, pX2, pY2, image):
    color = (0, 255, 0)  # BGR
    thickness = 1  # 0 、4、8
    n = pX1.size(0)
    for i in range(n):
        cv2.line(image, (int(pX1[i]), int(pY1[i])), (int(pX2[i]), int(pY2[i])), color, thickness)  # ,cv2.LINE_4)

def PlotDir2DCV(pAllPointList1, pAllOointList2, pPairPointList1, pPairPointList2, image1, image2, strName):
    pX1 = pAllPointList1[0]
    pY1 = pAllPointList1[1]
    pOrtX1 = pAllPointList1[2]
    pOrtY1 = pAllPointList1[3]

    pX2 = pAllOointList2[0]
    pY2 = pAllOointList2[1]
    pOrtX2 = pAllOointList2[2]
    pOrtY2 = pAllOointList2[3]

    height = image2.shape[0]
    width = image2.shape[1]

    mergeImage = np.concatenate((image1,image2),axis=1)
    mergeImage3 = np.ones(height * width * 2 * 3, dtype=np.uint8)
    mergeImage3 = mergeImage3.reshape(height, width*2, 3)
    mergeImage3[:, :, 0] = mergeImage
    mergeImage3[:, :, 1] = mergeImage
    mergeImage3[:, :, 2] = mergeImage

    PlotDirCircleCV(pX1, pY1, pOrtX1, pOrtY1, mergeImage3)
    PlotDirCircleCV(pX2+width, pY2, pOrtX2, pOrtY2, mergeImage3)

    pX1 = pPairPointList1[0].squeeze(1)
    pY1 = pPairPointList1[1].squeeze(1)
    pX2 = pPairPointList2[0].squeeze(1)
    pY2 = pPairPointList2[1].squeeze(1)
    Plot2DLineCV(pX1, pY1, pX2+width, pY2, mergeImage3)

    cv2.imwrite(strName, mergeImage3)

def testImageRegister(lambda3):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    name1 = "2002DB1A1_1"
    name2 = "2002DB1A1_2"
    pointAllList, pointPairList, image1, image2 = LoadImageData(name1, name2)
    height = image2.shape[0]
    width = image2.shape[1]
    PlotDir2DCV(pointAllList[0:4], pointAllList[4:8], pointPairList[0:4], pointPairList[4:8], image1, image2, "orgImage.png")

    pAllX2 = pointAllList[4].requires_grad_(True)
    pAllY2 = pointAllList[5].requires_grad_(True)
    pAllOrtX2 = pointAllList[6]
    pAllOrtY2 = pointAllList[7]

    pX1 = pointPairList[0]
    pY1 = pointPairList[1]
    pOrtX1 = pointPairList[2]
    pOrtY1 = pointPairList[3]
    pX2 = pointPairList[4]
    pY2 = pointPairList[5]
    pOrtX2 = pointPairList[6]
    pOrtY2 = pointPairList[7]
    u, v = GetTransFunUV(pX2, pY2, pOrtX2, pOrtY2, pX1, pY1, pOrtX1, pOrtY1, lambda3, device)

    pImageX, pImageY = getImageTransPoints(width, height, device)
    pTImageList = GetTransPoint(u, v, pImageX, pImageY, None, None)
    pTransImageX = pTImageList[0].to('cpu')
    pTransImageY = pTImageList[1].to('cpu')
    transImage2, mse = GetTransImagePixel(width, height, pTransImageX, pTransImageY, image1, image2)

    u2, v2 = GetTransFunUV(pX1, pY1, pOrtX1, pOrtY1, pX2, pY2, pOrtX2, pOrtY2, lambda3, device)
    u2 = u2.to('cpu')
    v2 = v2.to('cpu')
    pTAllX2 = GetTransPoint(u2, v2, pAllX2.unsqueeze(1), pAllY2.unsqueeze(1), pAllOrtX2.unsqueeze(1), pAllOrtY2.unsqueeze(1))
    pTX2 = GetTransPoint(u2, v2, pX2, pY2, None, None)
    name = "TransImage" + str(int(lambda3)) + ".png"
    PlotDir2DCV(pointAllList[0:4], pTAllX2, pointPairList[0:4], [pTX2[0].unsqueeze(1), pTX2[1].unsqueeze(1)], image1, transImage2, name) #"TransImage.png"

    print("mse:", mse)

    # toPIL = transforms.ToPILImage()  # 这个函数可以将张量转为PIL图片，由小数转为0-255之间的像素值
    # img_PIL = toPIL(transImage2).convert('RGB')  # 张量tensor转换为图片
    # img_PIL.save('pinn.jpg')
    # img_PIL.show()

if __name__ == '__main__':
    # a=math.cos(3.14159)
    # data = torch.rand(10)
    # b = round(float(data[0]))
    # transImage = torch.zeros(100, dtype=torch.uint8)
    a=125//20
    b=125%20

    print("传入参数的总长度为：", len(sys.argv))
    print("type:", type(sys.argv))
    print("function name:", sys.argv[0])
    try:
        print("第一个传入的参数为 lambda3:", sys.argv[1])
    except Exception as e:
        print("Input Error:", e)
    lambda3 = float(sys.argv[1])
    testImageRegister(lambda3)
    print("testImageRegister finish!")